from data_utils import BatchManager, get_dict
from model import Model
import tensorflow as tf
tf = tf.compat.v1

batch_size = 20
dict_file = 'data/prepare/dict.pkl'


def train():
    # ---------数据准备-----------
    train_manager = BatchManager(batch_size=20, name='train')

    # ---------读取字典-----------
    mappping_dict = get_dict(dict_file)

    # ---------搭建模型-----------
    model = Model(mappping_dict)

    init = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init)
        for i in range(10):
            for batch in train_manager.iter_batch(shuffle=True):
                _, loss = sess.run([model.train_op, model.cost], feed_dict={model.char_inputs: batch[0],
                                                                            model.bound_inputs: batch[1],
                                                                            model.flag_inputs: batch[2],
                                                                            model.radical_inputs: batch[3],
                                                                            model.pinyin_inputs: batch[4],
                                                                            model.targets: batch[5]})
                print('loss{}'.format(loss))


if __name__ == '__main__':
    train()
